import os
import base64
from io import BytesIO
from PIL import Image

import torch
import torch.nn as nn
import numpy as np

from fastapi import File, UploadFile, HTTPException, Path, APIRouter
from fastapi.responses import JSONResponse, FileResponse

from concurrent.futures import ThreadPoolExecutor
from typing import Optional, List
from torchvision import transforms

from utils.response_util import ResponseUtil
from ..fun.niqe import calculate_niqe_from_file_stream
from ..fun.flame_ocalization import detect_fire_in_image,detect_smoke_in_image
from skimage.metrics import structural_similarity as ssim


# -------------------- 定义网络结构 --------------------
class Network(nn.Module):
    def __init__(self, n_chan, chan_embed=48):
        super(Network, self).__init__()
        self.act = nn.LeakyReLU(negative_slope=0.2, inplace=True)
        self.conv1 = nn.Conv2d(n_chan, chan_embed, 3, padding=1)
        self.conv2 = nn.Conv2d(chan_embed, chan_embed, 3, padding=1)
        self.conv3 = nn.Conv2d(chan_embed, n_chan, 1)

    def forward(self, x):
        x = self.act(self.conv1(x))
        x = self.act(self.conv2(x))
        x = self.conv3(x)
        return x


# -------------------- MSE 计算函数 --------------------
def mse(gt: torch.Tensor, pred: torch.Tensor) -> torch.Tensor:
    loss = nn.MSELoss()
    return loss(gt, pred)


# -------------------- 加载模型 --------------------
def load_model(weights_path, n_chan=3, device='cuda'):
    model = Network(n_chan)
    model.load_state_dict(torch.load(weights_path, map_location=device))
    model.to(device)
    model.eval()
    return model


BASE_DIR = os.path.dirname(os.path.abspath(__file__))  # 获取当前脚本的绝对目录
MODEL_WEIGHTS_PATH = os.path.join(BASE_DIR, 'ckpt', 'Denoise_model.pth')
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
model = load_model(MODEL_WEIGHTS_PATH, n_chan=3, device=DEVICE)

# -------------------- 图像转换器 --------------------
transform = transforms.ToTensor()


# -------------------- 辅助函数：PIL -> Base64 --------------------
def pil_to_base64(img: Image.Image) -> str:
    buffered = BytesIO()
    img.save(buffered, format="PNG")
    img_bytes = buffered.getvalue()
    return base64.b64encode(img_bytes).decode('utf-8')


# -------------------- 辅助函数：去噪处理 --------------------
def denoise_image(noise_img: Image.Image, device: str = DEVICE) -> Image.Image:
    noi = transform(noise_img).unsqueeze(0).to(device)  # BCHW
    with torch.no_grad():
        pred = torch.clamp(noi - model(noi), 0, 1)
    denoi = (pred * 255).squeeze(0).permute(1, 2, 0).cpu().numpy().astype(np.uint8)
    return Image.fromarray(denoi)


# 为了代码简洁先注释掉相关检查
# -------------------- 检查文件是否允许 --------------------
# def allowed_file(filename: str) -> bool:
#     extension = filename.rsplit('.', 1)[1].lower() if '.' in filename else ''
#     return extension in ALLOWED_EXTENSIONS


# -------------------- 处理单个文件，计算NIQE分数 --------------------
def process_file(args):
    filename, file_bytes = args
    try:
        if filename:
            niqe_score = calculate_niqe_from_file_stream(file_bytes)
            niqe_score = round(niqe_score, 3) if niqe_score else None
            return filename, niqe_score
        return filename, None
    except Exception as e:
        print(f"Error processing file {filename}: {e}")
        return filename, None


# -------------------- 初始化全局线程池 --------------------
executor = ThreadPoolExecutor(max_workers=6)

bigScreenController = APIRouter(prefix='/big-screen')


# -------------------- 路由：去噪 --------------------
@bigScreenController.post("/denoise/{index}")
async def denoise_route(
        index: int = Path(..., description="示例: /denoise/0"),
        image1: Optional[UploadFile] = File(None),
        image2: Optional[UploadFile] = File(None),
        image3: Optional[UploadFile] = File(None),
):
    """
    接收 1~3 张图像，进行去噪处理，返回Base64格式结果。
    """
    try:
        files = [(image1, "image1"), (image2, "image2"), (image3, "image3")]
        denoised_images_b64 = []
        for file_obj, file_key in files:
            if file_obj is None:
                # 如果某张图没传，就跳过
                continue
            # if file_obj.filename == "":
            #     raise HTTPException(status_code=400, detail=f"No selected file for {file_key}")
            # if not allowed_file(file_obj.filename):
            #     raise HTTPException(status_code=400, detail=f"File type not allowed for {file_key}")

            # 读取并转换
            content = await file_obj.read()
            img = Image.open(BytesIO(content)).convert('RGB')
            denoised_img = denoise_image(img)
            denoised_b64 = pil_to_base64(denoised_img)
            denoised_images_b64.append(denoised_b64)

        return JSONResponse(content={
            "denoised_images": denoised_images_b64
        }, status_code=200)

    except Exception as e:
        print(f"Error in denoise_route: {e}")
        raise HTTPException(status_code=500, detail=str(e))


# -------------------- 路由：上传图片并计算 NIQE 分数 --------------------
@bigScreenController.post("/upload/{index}")
async def upload_images(
        index: int = Path(..., description="示例: /upload/0"),
        image1: UploadFile = File(...),
        image2: UploadFile = File(...),
        image3: UploadFile = File(...),
        image4: UploadFile = File(...),
        image5: UploadFile = File(...),
        image6: UploadFile = File(...)
):
    """
    接收6张图片，计算其 NIQE 分数。
    """
    try:
        # 这里的Key可自由调整
        file_objs = [image1, image2, image3, image4, image5, image6]
        results_list = []

        for file_obj in file_objs:
            if file_obj.filename == "":
                raise HTTPException(status_code=400, detail=f"No selected file for {file_obj.name}")
            # if not allowed_file(file_obj.filename):
            #     raise HTTPException(status_code=400, detail=f"File type not allowed for {file_obj.name}")

            # 读取文件并提交到线程池处理
            file_bytes = await file_obj.read()
            results_list.append((file_obj.filename, file_bytes))

        # 线程池并行处理
        futures = [executor.submit(process_file, item) for item in results_list]
        results = [f.result() for f in futures]

        response_data = {}
        score_list = []
        for filename, score in results:
            if score is not None:
                response_data[filename] = score
                score_list.append(score)
            else:
                response_data[filename] = "Error processing file"
                score_list.append(None)

        return JSONResponse(content={
            "data": response_data,
            "results": score_list
        }, status_code=200)

    except Exception as e:
        print(f"Error in upload_images: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@bigScreenController.post("/flame/{index}")
async def flame_images(
        index: int = Path(..., description="示例: /flame/0"),
        image1: UploadFile = File(...),
        image2: UploadFile = File(...),
        image3: UploadFile = File(...),
        image4: UploadFile = File(...),
        image5: UploadFile = File(...),
        image6: UploadFile = File(...)
):
    file_objs = [image1, image2, image3, image4, image5, image6]

    results = []

    for idx, file_obj in enumerate(file_objs, start=1):
        try:
            image_bytes = await file_obj.read()  # 读取上传文件的字节数据
            processed_bytes = detect_fire_in_image(image_bytes)

            # 将处理后的图像转换为 Base64 字符串，方便放进 JSON 返回
            processed_b64 = base64.b64encode(processed_bytes).decode("utf-8")

            # 也可以只存储并在前端自行转换，这里演示返回 Base64
            results.append({
                "filename": file_obj.filename or f"image_{idx}.jpg",
                "base64": processed_b64
            })
        except Exception as e:
            results.append({
                "filename": file_obj.filename or f"image_{idx}.jpg",
                "error": str(e)
            })

    return ResponseUtil.success(data=results)



@bigScreenController.post("/smoke/{index}")
async def somke_images(
        index: int = Path(..., description="示例: /flame/0"),
        image1: UploadFile = File(...),
        image2: UploadFile = File(...),
        image3: UploadFile = File(...),
        image4: UploadFile = File(...),
        image5: UploadFile = File(...),
        image6: UploadFile = File(...)
):
    file_objs = [image1, image2, image3, image4, image5, image6]

    results = []

    for idx, file_obj in enumerate(file_objs, start=1):
        try:
            image_bytes = await file_obj.read()  # 读取上传文件的字节数据
            processed_bytes = detect_smoke_in_image(image_bytes)

            # 将处理后的图像转换为 Base64 字符串，方便放进 JSON 返回
            processed_b64 = base64.b64encode(processed_bytes).decode("utf-8")

            # 也可以只存储并在前端自行转换，这里演示返回 Base64
            results.append({
                "filename": file_obj.filename or f"image_{idx}.jpg",
                "base64": processed_b64
            })
        except Exception as e:
            results.append({
                "filename": file_obj.filename or f"image_{idx}.jpg",
                "error": str(e)
            })

    return ResponseUtil.success(data=results)